from typing import List, Dict

import pandas as pd
from fastapi import APIRouter, Body
from pydantic import BaseModel

from app.utils.ai_util import load_data_csv, load_emb_csv, emb_do_query
from dataclasses import dataclass


router = APIRouter()


@dataclass(frozen=True)
class EmbResource:
    df_209: pd.DataFrame
    df_601: pd.DataFrame
    df_209_emb: pd.DataFrame
    df_601_emb: pd.DataFrame


def resource_init():
    df_209, df_601 = load_data_csv()
    df_209_emb, df_601_emb = load_emb_csv()
    return EmbResource(df_209=df_209, df_601=df_601, df_209_emb=df_209_emb, df_601_emb=df_601_emb)


resource_data: EmbResource = resource_init()


class EmbReq(BaseModel):
    emp_type: str
    query: str
    history: List[Dict] = None
    temperature: int = 0.1
    top_n: int = 4


class EmbRes(BaseModel):
    content: str


@router.post("/query", response_model=EmbRes)
def ai_emb_query(
    emb_req: EmbReq = Body(example=dict(emp_type="emp / store_emp", query="查询内容"))
):

    df, df_emb = None, None
    if emb_req.emp_type == "emp":
        df, df_emb = resource_data.df_209, resource_data.df_209_emb
    elif emb_req.emp_type == "store_emp":
        df, df_emb = resource_data.df_601, resource_data.df_601_emb
    else:
        return EmbRes(content="emp_type 不对：{}".format(emb_req.emp_type))
    content = emb_do_query(
        query=emb_req.query, df=df, df_emb=df_emb, history=emb_req.history, temperature=emb_req.temperature
    )
    return EmbRes(content=content)
